import torch.nn as nn
from spikingjelly.clock_driven import neuron
import torch.nn.functional as F
class Net(nn.Module):
    def __init__(self, tau=2.0, v_threshold=1.0, v_reset=0.0):
        super().__init__()
        # 网络结构，简单的双层全连接网络，每一层之后都是LIF神经元
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 10, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset),
        )

    def forward(self, x):
        return self.fc(x)
# Snet为SNN模型，直接修改其中内容即可
class Snet(nn.Module):
    def __init__(self, tau=2.0, v_threshold=1.0, v_reset=0.0):
        super().__init__()
        # 网络结构，简单的双层全连接网络，每一层之后都是LIF神经元
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 14 * 14, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset),
            nn.Linear(14 * 14, 10, bias=False),
            neuron.LIFNode(tau=tau, v_threshold=v_threshold, v_reset=v_reset)
        )
    def forward(self, x):
        return self.fc(x)

class CNN(nn.Module):
    # 定义一个简单的网络
    # LeNet -5
    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5)
            self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
            self.fc1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
            self.fc2 = nn.Linear(in_features=120, out_features=84)
            self.fc3 = nn.Linear(in_features=84, out_features=10)

        def forward(self, x):
            x = self.pool1(F.relu(self.conv1(x)))
            x = self.pool1(F.relu(self.conv2(x)))
            x = x.view(-1, 16 * 5 * 5)  # reshape tensor
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            return x
